'''
返回一个train和test的Dataloader
'''
'''
输入 xdata ydata均为nparray
'''
import torch.utils.data as Data
import torch
import matplotlib.pyplot as plt
from prettytable import PrettyTable
import numpy as np

from sklearn.model_selection import train_test_split

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

def draw_loss(Loss_list,epoch):
    # 我这里迭代了200次，所以x的取值范围为(0，200)，然后再将每次相对应的准确率以及损失率附在x上
    plt.cla()
    x1 = range(1, epoch+1)
    y1 = Loss_list
    plt.title('Train loss vs. epoches', fontsize=20)
    plt.plot(x1, y1, '.-')
    plt.xlabel('epoches', fontsize=20)
    plt.ylabel('Train loss', fontsize=20)
    plt.grid()
    plt.savefig("./Train_loss.png")
    plt.show()


def draw_fig(list_data,name,epoch):
    # 我这里迭代了200次，所以x的取值范围为(0，200)，然后再将每次相对应的准确率以及损失率附在x上
    x1 = range(1, epoch+1)
    y1 = list_data
    if name=="loss":
        plt.cla()
        plt.title('Train loss vs. epoch', fontsize=15)
        plt.plot(x1, y1, '.-')
        plt.xlabel('epoch', fontsize=20)
        plt.ylabel('Train loss', fontsize=20)
        plt.grid()
        plt.savefig("./Train_loss.png")
        plt.show()
    elif name =="acc":
        plt.cla()
        plt.title('Test Accuracy vs. Epoch', fontsize=15)
        plt.plot(x1, y1, '.-')
        plt.xlabel('Epoch', fontsize=20)
        plt.ylabel('Test Acc', fontsize=20)
        plt.grid()
        plt.savefig("./Val _accuracy.png")
        plt.show()

def datahandle(xdata, ydata, batchsize,randomstate,splitsize=0.3,isshuffle=False):
    x, y = torch.tensor(xdata).to(torch.float32), torch.tensor(ydata).to(torch.float32)
    x = x.to(device)
    y = y.to(device)
    X_train,X_test, y_train, y_test = train_test_split(x,y,test_size=splitsize, random_state=randomstate)

    print("test_size:", len(X_test), "train_size:", len(X_train))

    train_loader = Data.DataLoader(
        dataset=Data.TensorDataset(X_train, y_train),
        batch_size=batchsize,
        shuffle=isshuffle,
    )

    val_loader = Data.DataLoader(
        dataset=Data.TensorDataset(X_test, y_test),
        batch_size=batchsize,
        shuffle=isshuffle,
    )
    return train_loader,val_loader


def returndataaloader(xtrain,xtest,ytrain,ytest,batchsize,isshuffle = False):
    x1, y1 = torch.tensor(xtrain).to(torch.float32), torch.tensor(ytrain).to(torch.float32)
    x2, y2 = torch.tensor(xtest).to(torch.float32), torch.tensor(ytest).to(torch.float32)
    x1,x2 = x1.to(device),x2.to(device)
    y1,y2 = y1.to(device),y2.to(device)
    train_loader = Data.DataLoader(
        dataset=Data.TensorDataset(x1, y1),
        batch_size=batchsize,
        shuffle=isshuffle,
    )

    val_loader = Data.DataLoader(
        dataset=Data.TensorDataset(x2, y2),
        batch_size=batchsize,
        shuffle=isshuffle,
    )

    return train_loader, val_loader



class ConfusionMatrix(object):

    def __init__(self, num_classes: int, labels: list):
        self.matrix = np.zeros((num_classes, num_classes))  # 初始化混淆矩阵，元素都为0
        self.num_classes = num_classes  # 类别数量，本例数据集类别为5
        self.labels = labels  # 类别标签

    def update(self, preds, labels):
        for p, t in zip(preds, labels):  # pred为预测结果，labels为真实标签
            self.matrix[p, t] += 1  # 根据预测结果和真实标签的值统计数量，在混淆矩阵相应位置+1

    def summary(self):  # 计算指标函数
        # calculate accuracy
        sum_TP = 0
        n = np.sum(self.matrix)
        for i in range(self.num_classes):
            sum_TP += self.matrix[i, i]  # 混淆矩阵对角线的元素之和，也就是分类正确的数量
        acc = sum_TP / n  # 总体准确率
        print("the model accuracy is ", acc)

        # kappa
        sum_po = 0
        sum_pe = 0
        for i in range(len(self.matrix[0])):
            sum_po += self.matrix[i][i]
            row = np.sum(self.matrix[i, :])
            col = np.sum(self.matrix[:, i])
            sum_pe += row * col
        po = sum_po / n
        pe = sum_pe / (n * n)
        # print(po, pe)
        kappa = round((po - pe) / (1 - pe), 3)
        # print("the model kappa is ", kappa)

        # precision, recall, specificity
        table = PrettyTable()  # 创建一个表格
        table.field_names = ["", "Precision", "Recall", "Specificity"]
        for i in range(self.num_classes):  # 精确度、召回率、特异度的计算
            TP = self.matrix[i, i]
            FP = np.sum(self.matrix[i, :]) - TP
            FN = np.sum(self.matrix[:, i]) - TP
            TN = np.sum(self.matrix) - TP - FP - FN

            Precision = round(TP / (TP + FP), 3) if TP + FP != 0 else 0.
            Recall = round(TP / (TP + FN), 3) if TP + FN != 0 else 0.  # 每一类准确度
            Specificity = round(TN / (TN + FP), 3) if TN + FP != 0 else 0.

            table.add_row([self.labels[i], Precision, Recall, Specificity])
        print(table)
        return str(acc)

    def plot(self):  # 绘制混淆矩阵
        matrix = self.matrix
        print(matrix)
        plt.imshow(matrix, cmap=plt.cm.Blues)

        # 设置x轴坐标label
        plt.xticks(range(self.num_classes), self.labels, rotation=45)
        # 设置y轴坐标label
        plt.yticks(range(self.num_classes), self.labels)
        # 显示colorbar
        plt.colorbar()
        plt.xlabel('True Labels')
        plt.ylabel('Predicted Labels')
        plt.title('Confusion matrix (acc=' + self.summary() + ')')

        # 在图中标注数量/概率信息
        thresh = matrix.max() / 2
        for x in range(self.num_classes):
            for y in range(self.num_classes):
                # 注意这里的matrix[y, x]不是matrix[x, y]
                info = int(matrix[y, x])
                plt.text(x, y, info,
                         verticalalignment='center',
                         horizontalalignment='center',
                         color="white" if info > thresh else "black")
        plt.tight_layout()
        plt.show()

def data_process(xfile,yfile,n=5000):
    sampleindex = np.sort(np.random.choice([i for i in range(xfile.shape[0])],n,replace=False))
    X = xfile[sampleindex,:]
    Y = yfile[sampleindex,:]
    print("X_SHAPE:",X.shape," Y_SHAPE:",Y.shape)
    return X,Y

def classAcc(pre1,lab1):
    pre1 = pre1.numpy().reshape(1,-1)
    lab1 = lab1.numpy().reshape(1,-1)
    # print(pre1.shape,lab1.shape)
    return np.sum(pre1==lab1)/lab1.shape[1]
def fprloss(pre1,lab1):
    fprcount = 0
    tprcount = 0
    for index1,val1 in enumerate(pre1):
        if (val1==0 and lab1[index1]==1):
            fprcount+=1
        if (val1==1 and lab1[index1]==0):
            tprcount+=1
    if(len(lab1.nonzero())==0):
        return tprcount/len(lab1)
    return fprcount/len(lab1.nonzero())


def option2svmy(y):
    tmp = []
    for i in y:
        for index,v in enumerate(i):
            if v == 1:
                tmp.append(index)
                break
    return np.array(tmp)